#pragma once

#include "util/hashmap/hash_algorithm.h"
#include "util/hashmap/hash_bucket.h"
#include "util/hashmap/hash_util.h"
#include "util/pool/pool.h"

namespace hawking {
namespace indexlib {

template <typename Key, typename Pool = Pool, typename HashFunc = KeyHash<Key>,
          typename Comp = KeyEqual<Key>, int Sparsity = 2>
class HashSet {
    using Typed = HashSet<Key, Pool, HashFunc, Comp>;
    HashSet(const Typed&) = delete;
    Typed& operator=(const Typed&) = delete;

public:
    typedef Key KeyType;
    typedef Pool PoolType;
    typedef HashBucket<KeyType, PoolType> BucketType;
    typedef typename BucketType::Block BlockType;
    typedef HashSet<Key, Pool, HashFunc, Comp> HashSetType;

    static const size_t DEFAULT_HASHSET_POOL_SIZE = 10 * 1024 * 1024; // 10M

public:
    class Iterator
    {
    public:
        Iterator() :
            map_self_(nullptr), buckets_(nullptr),
            current_block_(nullptr), block_item_pos_(0) {}

        Iterator(const HashSet* m) :
            map_self_(m), buckets_(map_self_->bucket_),
            current_block_(nullptr), block_item_pos_(0) {
            Reset();
        }

        void Reset() {
            current_block_ = buckets_->GetFirstBlock();
            while (current_block_) {
                Bitmap& bitmap = current_block_->bitmap;
                block_item_pos_ = 0;
                while (true) {
                    if (block_item_pos_ >= current_block_->size) {
                        break;
                    }
                    if (bitmap.Test(block_item_pos_)) {
                        return;
                    }
                    ++block_item_pos_;
                }
                current_block_ = current_block_->next;
            }
        }

        bool HasNext() const {
            return current_block_ && current_block_->bitmap.Test(block_item_pos_);
        }

        KeyType& Next() {
            KeyType& key = current_block_->begin[block_item_pos_++];
            while (current_block_) {
                Bitmap& bitmap = current_block_->bitmap;
                while (true) {
                    if (block_item_pos_ >= current_block_->size) {
                        break;
                    }
                    if (bitmap.Test(block_item_pos_)) {
                        return key;
                    }
                    ++block_item_pos_;
                }

                current_block_ = current_block_->next;
                block_item_pos_ = 0;
            }

            return key;
        }

    private:
        const HashSet* map_self_ = nullptr;
        BucketType* buckets_ = nullptr;
        BlockType* current_block_ = nullptr;
        size_t block_item_pos_ = 0;
    };

public:
    HashSet(PoolType* pool, size_t init_size);
    HashSet(size_t init_size);
    HashSet(size_t init_size, size_t pool_size);
    ~HashSet();

public:
    bool Exist(const KeyType& x) const;

    // the same key may have duplicated entries
    void Insert(const KeyType& key);
    bool operator==(const HashSetType& otherMap) const;
    Iterator CreateIterator() const { return Iterator(this); }

    static size_t EstimateFirstBucketSize(size_t init_size);
    static int64_t EstimateNeededMemory(int64_t keyCount);

    size_t Size() const { return element_count_; }
    void Clear();

private:
    size_t Threshold_();
    KeyType* Lookup_(const KeyType& x) const;
    void FindAndInsertInLastBucket_(const KeyType& x);
    bool IsValid_() const;

    HashFunc hasher_;
    Comp comparator_;
    size_t element_count_ = 0;
    BucketType* bucket_ = nullptr;

    PoolType* mem_pool_ = nullptr;
    bool own_pool_ = false;
};

/////////////////////////////////////////////////////////
// Implementation
template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
HashSet<Key, Pool, HashFunc, Comp, Sparsity>::HashSet(PoolType* pool, size_t init_size)
    : element_count_(0)
    , bucket_(new BucketType(pool, init_size * Sparsity))
    , mem_pool_(pool)
    , own_pool_(false) {
    assert(mem_pool_ != nullptr);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
HashSet<Key, Pool, HashFunc, Comp, Sparsity>::HashSet(size_t init_size)
    : element_count_(0)
    , mem_pool_(new PoolType(DEFAULT_HASHSET_POOL_SIZE))
    , own_pool_(true) {
    bucket_ = new BucketType(mem_pool_, init_size * Sparsity);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
HashSet<Key, Pool, HashFunc, Comp, Sparsity>::HashSet(
    size_t init_size, size_t pool_size)
    : element_count_(0)
    , mem_pool_(new PoolType(pool_size))
    , own_pool_(true) {
    bucket_ = new BucketType(mem_pool_, init_size * Sparsity);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
HashSet<Key, Pool, HashFunc, Comp, Sparsity>::~HashSet() {
    Clear();

    delete bucket_;
    bucket_ = nullptr;
    if (own_pool_ && mem_pool_) {
        mem_pool_->Release();
        delete mem_pool_;
        mem_pool_ = nullptr;
    }
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline bool HashSet<Key, Pool, HashFunc, Comp, Sparsity>::Exist(const KeyType& x) const {
    assert(IsValid_());
    KeyType* p = Lookup_(x);
    return p != nullptr;
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline void HashSet<Key, Pool, HashFunc, Comp, Sparsity>::Insert(const KeyType& x) {
    assert(IsValid_());

    if (Lookup_(x)) {
        return;
    }

    if (element_count_ >= Threshold_()) {
        bucket_->AddBlock();
    }

    FindAndInsertInLastBucket_(x);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
void HashSet<Key, Pool, HashFunc, Comp, Sparsity>::Clear() {
    bucket_->Clear();
    element_count_ = 0;
    if (own_pool_ && mem_pool_) {
        mem_pool_->Reset();
    }
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline size_t HashSet<Key, Pool, HashFunc, Comp, Sparsity>::Threshold_() {
    return bucket_->Size() / Sparsity;
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline typename HashSet<Key, Pool, HashFunc, Comp, Sparsity>::KeyType*
HashSet<Key, Pool, HashFunc, Comp, Sparsity>::Lookup_(const KeyType& x) const {
    BlockType* buff = bucket_->GetLastBlock();
    while (buff) {
        Bitmap& bitmap = buff->bitmap;
        size_t pos = hasher_(x) % buff->size;
        while (bitmap.Test(pos)) {
            KeyType* p = buff->begin + pos;
            if (comparator_(*p, x)) {
                return p;
            }
            if (++pos >= buff->size) {
                pos = 0;
            }
        }
        buff = buff->prev;
    }

    return nullptr;
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline void HashSet<Key, Pool, HashFunc, Comp, Sparsity>::FindAndInsertInLastBucket_(
    const KeyType& x) {
    BlockType* buff = bucket_->GetLastBlock();
    Bitmap& bitmap = buff->bitmap;
    size_t pos = hasher_(x) % buff->size;

    KeyType* p = buff->begin + pos;
    bool isFound = false;
    while (bitmap.Test(pos)) {
        if (comparator_(*p, x)) {
            isFound = true;
            break;
        }
        if (++pos >= buff->size) {
            pos = 0;
        }
        p = buff->begin + pos;
    }
    if (!isFound) {
        element_count_++;
        bucket_->GetLastBlock()->capacity++;
    }

    *p = x;

    __asm__ __volatile__("" ::: "memory");
    bitmap.Set(pos);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
size_t HashSet<Key, Pool, HashFunc, Comp, Sparsity>::EstimateFirstBucketSize(
    size_t init_size) {
    size_t allocElemSize = BucketType::GetNextBlockSize(init_size * Sparsity);
    int32_t slotCount = Bitmap::GetSlotCount(allocElemSize);
    return allocElemSize * sizeof(KeyType) +
        sizeof(BlockType) + slotCount * sizeof(uint32_t);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
int64_t HashSet<Key, Pool, HashFunc, Comp, Sparsity>::EstimateNeededMemory(
    int64_t keyCount) {
    int64_t memNeed = 0;
    int64_t allocElemSize = 0;
    int64_t capacity = 0;
    while (keyCount > (capacity / Sparsity)) {
        allocElemSize = BucketType::GetNextBlockSize(allocElemSize);
        capacity += allocElemSize;
        int32_t slotCount = Bitmap::GetSlotCount(allocElemSize);
        memNeed +=
            allocElemSize * sizeof(KeyType) +
                sizeof(BlockType) + slotCount * sizeof(uint32_t);
    }

    return memNeed;
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline bool HashSet<Key, Pool, HashFunc, Comp, Sparsity>::IsValid_() const {
    return (element_count_ >= 0) && (element_count_ <= (size_t)bucket_->Size() / Sparsity);
}

template <typename Key, typename Pool, typename HashFunc, typename Comp, int Sparsity>
inline bool HashSet<Key, Pool, HashFunc, Comp, Sparsity>::operator==(
    const HashSetType& otherMap) const {
    if (element_count_ != otherMap.element_count_) {
        return false;
    }

    BlockType* buff = bucket_->GetLastBlock();
    BlockType* otherBuff = otherMap.bucket_->GetLastBlock();

    while (buff != nullptr && otherBuff != nullptr) {
        // compare current block
        if (buff->size != otherBuff->size) {
            return false;
        }
        Bitmap& bitmap = buff->bitmap;
        Bitmap& otherBitmap = otherBuff->bitmap;
        if (bitmap != otherBitmap) {
            return false;
        }
        for (size_t i = 0; i < buff->size; ++i) {
            if (!bitmap.Test(i)) {
                continue;
            }
            if (buff->begin[i] != otherBuff->begin[i]) {
                return false;
            }
        }
        buff = buff->prev;
        otherBuff = otherBuff->prev;
    }

    if ((buff == nullptr && otherBuff != nullptr) ||
        (buff != nullptr && otherBuff == nullptr)) {
        return false;
    }
    return true;
}

}
}